#!/usr/bin/env python3
"""
数据库操作模块
支持PostgreSQL和MySQL数据库连接和查询
"""

import logging
from datetime import datetime
from typing import List, Dict, Any
from config import Config

logger = logging.getLogger(__name__)

class ZabbixDatabase:
    """Zabbix数据库查询类 - 支持PostgreSQL和MySQL"""
    
    def __init__(self):
        self.connection = None
        self.db_config = Config.get_db_config()
        self.db_type = self.db_config['type']
        
    def connect(self):
        """连接数据库 - 支持PostgreSQL和MySQL"""
        try:
            if self.db_type == 'postgresql':
                self._connect_postgresql()
            elif self.db_type == 'mysql':
                self._connect_mysql()
            else:
                raise ValueError(f"不支持的数据库类型: {self.db_type}。支持的类型: postgresql, mysql")
                
        except ImportError as e:
            self._handle_import_error(e)
            raise
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise
    
    def _connect_postgresql(self):
        """连接PostgreSQL数据库"""
        import psycopg2
        self.connection = psycopg2.connect(
            host=self.db_config['host'],
            port=self.db_config['port'],
            user=self.db_config['user'],
            password=self.db_config['password'],
            database=self.db_config['database'],
            connect_timeout=10
        )
        logger.info(f"成功连接到 PostgreSQL 数据库: {self.db_config['host']}:{self.db_config['port']}")
    
    def _connect_mysql(self):
        """连接MySQL数据库"""
        import pymysql
        self.connection = pymysql.connect(
            host=self.db_config['host'],
            port=self.db_config['port'],
            user=self.db_config['user'],
            password=self.db_config['password'],
            database=self.db_config['database'],
            charset='utf8mb4',
            connect_timeout=10,
            cursorclass=pymysql.cursors.DictCursor
        )
        logger.info(f"成功连接到 MySQL 数据库: {self.db_config['host']}:{self.db_config['port']}")
    
    def _handle_import_error(self, error):
        """处理驱动导入错误"""
        if 'psycopg2' in str(error):
            logger.error("PostgreSQL驱动未安装，请运行: pip install psycopg2-binary")
        elif 'pymysql' in str(error):
            logger.error("MySQL驱动未安装，请运行: pip install pymysql")
    
    def disconnect(self):
        """断开数据库连接"""
        if self.connection:
            self.connection.close()
            logger.info("数据库连接已关闭")
    
    def get_current_alerts(self) -> List[Dict[str, Any]]:
        """获取当前告警信息 - 支持PostgreSQL和MySQL和多版本Zabbix"""
        try:
            # 先检测Zabbix版本和表结构
            zabbix_version = self._detect_zabbix_version()
            
            if zabbix_version >= 6.0:
                # Zabbix 6.0+ 使用problem表
                sql = self._get_problem_table_sql()
                alerts = self._execute_query(sql)
                if not alerts:
                    logger.warning("problem表查询未返回结果，尝试events表")
                    sql = self._get_events_table_sql()
                    alerts = self._execute_query(sql)
            else:
                # Zabbix 5.x 及以下使用events表
                sql = self._get_events_table_sql()
                alerts = self._execute_query(sql)
            
            # 记录最终结果含版本信息
            logger.info(f"获取到 {len(alerts)} 条当前告警 (Zabbix {zabbix_version}, {self.db_type.upper()})") 
            return alerts
            
        except Exception as e:
            logger.error(f"查询告警失败: {e}")
            logger.info("尝试使用简化查询")
            return self._get_alerts_fallback()
    
    def _execute_query(self, sql: str) -> List[Dict[str, Any]]:
        """执行SQL查询并返回结果"""
        try:
            cursor = self.connection.cursor()
            cursor.execute(sql)
            
            alerts = []
            
            if self.db_type == 'postgresql':
                alerts = self._process_postgresql_results(cursor)
            elif self.db_type == 'mysql':
                alerts = self._process_mysql_results(cursor)
            
            cursor.close()
            return alerts
            
        except Exception as e:
            logger.error(f"SQL执行失败: {e}")
            # PostgreSQL需要重置事务
            if self.db_type == 'postgresql':
                try:
                    self.connection.rollback()
                except Exception:
                    pass
            return []
    
    def _get_problem_table_sql(self) -> str:
        """获取Zabbix 6.0+的problem表查询SQL - 适配实际数据库结构"""
        return """
        SELECT 
            p.eventid as problemid,
            p.clock as problem_clock,
            p.r_clock as recovery_clock,
            p.name as problem_name,
            p.severity,
            t.description as trigger_description,
            h.host,
            h.name as host_name,
            p.name as event_name,
            '' as opdata,
            p.source,
            p.objectid
        FROM problem p
        LEFT JOIN triggers t ON p.objectid = t.triggerid
        LEFT JOIN functions f ON t.triggerid = f.triggerid
        LEFT JOIN items i ON f.itemid = i.itemid
        LEFT JOIN hosts h ON i.hostid = h.hostid
        WHERE p.r_clock = 0
        ORDER BY p.severity DESC, p.clock DESC
        LIMIT 100
        """
    
    def _get_events_table_sql(self) -> str:
        """获取Zabbix 5.x及以下的events表查询SQL - 适配实际数据库结构"""
        return """
        SELECT 
            e.eventid as problemid,
            e.clock as problem_clock,
            0 as recovery_clock,
            COALESCE(e.name, t.description) as problem_name,
            e.severity,
            t.description as trigger_description,
            h.host,
            h.name as host_name,
            e.name as event_name,
            '' as opdata,
            e.source,
            e.objectid
        FROM events e
        LEFT JOIN triggers t ON e.objectid = t.triggerid
        LEFT JOIN functions f ON t.triggerid = f.triggerid
        LEFT JOIN items i ON f.itemid = i.itemid
        LEFT JOIN hosts h ON i.hostid = h.hostid
        WHERE e.source = 0 
          AND e.value = 1
          AND NOT EXISTS (
              SELECT 1 FROM events e2 
              WHERE e2.source = 0 
                AND e2.object = e.object 
                AND e2.objectid = e.objectid 
                AND e2.value = 0 
                AND e2.clock > e.clock
          )
        ORDER BY e.severity DESC, e.clock DESC
        LIMIT 100
        """
    
    def _detect_zabbix_version(self) -> float:
        """检测Zabbix数据库版本 - 适配实际dbversion表结构"""
        try:
            cursor = self.connection.cursor()
            
            # 先尝试查看表是否存在
            cursor.execute("""
                SELECT table_name FROM information_schema.tables 
                WHERE table_name = 'dbversion'
            """)
            
            if cursor.fetchone():
                # 根据您的表结构，使用mandatory字段
                cursor.execute("SELECT mandatory FROM dbversion ORDER BY dbversionid DESC LIMIT 1")
                result = cursor.fetchone()
                if result:
                    if self.db_type == 'postgresql':
                        version_int = int(result[0]) if result[0] else 6000000
                    else:
                        version_int = int(result['mandatory'] if 'mandatory' in result else result[0]) if result else 6000000
                    
                    # 转换版本号 (6000000 = 6.0.0)
                    major = version_int // 1000000
                    minor = (version_int % 1000000) // 10000
                    version = float(f"{major}.{minor}")
                    cursor.close()
                    logger.info(f"检测到Zabbix版本: {version} (数据库值: {version_int})")
                    return version
            
            cursor.close()
            logger.warning("未找到dbversion表，默认作为6.0")
            return 6.0
            
        except Exception as e:
            logger.warning(f"无法检测Zabbix版本: {e}，默认作为6.0")
            # PostgreSQL需要重置事务
            if self.db_type == 'postgresql':
                try:
                    self.connection.rollback()
                except Exception:
                    pass
            return 6.0
    
    def _get_alerts_fallback(self) -> List[Dict[str, Any]]:
        """简化的后备查询方案 - 适配实际数据库结构"""
        try:
            # 最简单的triggers表查询
            sql = """
            SELECT 
                t.triggerid as problemid,
                EXTRACT(EPOCH FROM NOW())::integer as problem_clock,
                0 as recovery_clock,
                t.description as problem_name,
                t.priority as severity,
                t.description as trigger_description,
                h.host,
                h.name as host_name,
                'Trigger Alert' as event_name,
                '' as opdata,
                0 as source,
                t.triggerid as objectid
            FROM triggers t
            LEFT JOIN functions f ON t.triggerid = f.triggerid
            LEFT JOIN items i ON f.itemid = i.itemid
            LEFT JOIN hosts h ON i.hostid = h.hostid
            WHERE t.status = 0 AND t.value = 1
            ORDER BY t.priority DESC
            LIMIT 50
            """
            
            return self._execute_query(sql)
            
        except Exception as e:
            logger.error(f"后备查询也失败: {e}")
            return []
    
    def _process_postgresql_results(self, cursor) -> List[Dict[str, Any]]:
        """处理PostgreSQL查询结果"""
        columns = [desc[0] for desc in cursor.description]
        alerts = []
        
        for row in cursor.fetchall():
            alert = dict(zip(columns, row))
            self._process_alert_data(alert)
            alerts.append(alert)
        
        return alerts
    
    def _process_mysql_results(self, cursor) -> List[Dict[str, Any]]:
        """处理MySQL查询结果"""
        alerts = []
        
        for row in cursor.fetchall():
            alert = row  # 已经是字典格式
            self._process_alert_data(alert)
            alerts.append(alert)
        
        return alerts
    
    def _process_alert_data(self, alert: Dict[str, Any]):
        """处理告警数据通用逻辑 - 修复None值问题"""
        # 转换时间戳
        if alert.get('problem_clock'):
            try:
                alert['problem_time'] = datetime.fromtimestamp(alert['problem_clock']).strftime('%Y-%m-%d %H:%M:%S')
            except (ValueError, TypeError):
                alert['problem_time'] = '未知'
        else:
            alert['problem_time'] = '未知'
        
        # 处理可能的None值
        alert['host_name'] = alert.get('host_name') or alert.get('host') or '未知主机'
        alert['problem_name'] = alert.get('problem_name') or '未知告警'
        
        # 修复严重程度的None值问题
        if alert.get('severity') is None:
            alert['severity'] = 0  # 默认为未分类
        else:
            try:
                alert['severity'] = int(alert['severity'])
            except (ValueError, TypeError):
                alert['severity'] = 0
        
        # 确保其他字段不为None
        alert['trigger_description'] = alert.get('trigger_description') or alert.get('problem_name') or '无描述'
        alert['event_name'] = alert.get('event_name') or 'Zabbix Alert'
        alert['source'] = alert.get('source') or 0
        alert['objectid'] = alert.get('objectid') or 0